'''
文件说明
用途：根据检测结果txt文件内容在原图上绘制bounding box

检测结果格式：yolo检测结果
label_dir
|-----image1.txt
|-----image2.txt
|-----image3.txt

其中txt中格式内容为yolo检测结果的转化
# 每一行为对应图片检测到的bounding box
x_top_left, y_top_left, width, height, conf
x_top_left, y_top_left, width, height, conf
...

参考：
github：https://github.com/waittim/draw-YOLO-box/blob/main/draw_box.py
author：周小龙
date：24-12-1
'''
import cv2
from pathlib import Path
import random
from tqdm import tqdm
import warnings

def get_boxes(labels_path: Path):
    boxes = []
    with open(labels_path, 'r', encoding='utf-8') as file:
        for line in file:
            # 使用map函数将字符串转换为相应的数据类型
            parts = line.strip().split(',')
            # 将类别的字符串转换为整数，其余的转换为浮点数
            x_top_left = int(float(parts[0]))
            y_top_left = int(float(parts[1]))
            width = int(float(parts[2]))
            height = int(float(parts[3]))
            conf = float(parts[4])
            # 计算矩形的右下角坐标
            x_bottom_right = x_top_left + width
            y_bottom_right = y_top_left + height
            # 将转换后的数据添加到列表中
            boxes.append([x_top_left, y_top_left, x_bottom_right, y_bottom_right, conf])
    # print(boxes) # 展示bounding box
    return boxes

def draw(boxes, image, color=(0,0,255), line_thickness=None):
    for box in boxes:
        # Plots one bounding box on image img
        tl = line_thickness or round(
            0.002 * (image.shape[0] + image.shape[1]) / 2) + 1  # line/font thickness
        color = color or [random.randint(0, 255) for _ in range(3)]
        c1, c2, conf = (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), box[4]
        cv2.rectangle(image, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
        label = f"person {conf:.2f}"
        if label:
            tf = max(tl - 1, 1)  # font thickness
            t_size = cv2.getTextSize(label, 0, fontScale=tl / 4, thickness=tf)[0]
            c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
            cv2.rectangle(image, c1, c2, color, -1, cv2.LINE_AA)  # filled
            cv2.putText(image, label, (c1[0], c1[1] - 2), 0, tl / 4,
                        [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)

def _draw_box_old(labels_dir:Path, image_dir:Path, save_dir:Path):
    '''
    根据检测结果在原图绘制bounding box
    :param labels_dir: yolo检测结果对应的目录
    :param image_dir: 对应图片的目录
    :param save_dir: 绘制图片将要保存的目录
    :return:
    '''
    # 找出所有label
    warnings.warn("此方法已废弃，不推荐使用", DeprecationWarning)
    labels = [file for file in labels_dir.rglob("*") if file.is_file() if file.suffix in ['.txt']]
    for label in tqdm(labels):
        # 读取box
        boxes = get_boxes(label)
        # 读取对应图片
        img_path = Path(image_dir, label.stem).with_suffix('.jpg')
        # print(img_path) # 展示图片
        img = cv2.imread(img_path)
        draw(boxes, img)
        if not save_dir.exists():
            save_dir.mkdir(parents=True)
        cv2.imwrite(Path(save_dir, img_path.name), img)

def draw_box(labels_dir:Path, image_dir:Path, save_dir:Path, suffix = ['.jpg']):
    '''
    重构
    根据检测结果在原图绘制bounding box
    :param labels_dir: yolo检测结果对应的目录
    :param image_dir: 对应图片的目录
    :param save_dir: 绘制图片将要保存的目录
    :return:
    '''
    # 使用map存储标签
    # 读取图片
    imgs = [file for file in image_dir.rglob("*") if file.is_file() if file.suffix in suffix]

    # 找出所有label
    labels = [file for file in labels_dir.rglob("*") if file.is_file() if file.suffix in ['.txt']]
    # imgs_labels = {k:v for k,v in zip([label.stem for label in labels], labels)}
    imgs_labels = dict(zip([label.stem for label in labels], labels))

    for img_path in tqdm(imgs):
        # 读取图片
        img = cv2.imread(img_path)

        # 读取bounding box
        if img_path.stem in imgs_labels:
            label_path = imgs_labels.get(img_path.stem)
            boxes = get_boxes(label_path)
            draw(boxes, img)
        if not save_dir.exists():
            save_dir.mkdir(parents=True)
        cv2.imwrite(Path(save_dir, img_path.name), img)


def draw_single_img(img_path:Path, label_path:Path, save_dir:Path):
    '''
    绘制单一图片
    :param img_path:
    :param label_path:
    :param save_dir:
    :return:
    '''
    # 读取图片
    for i in tqdm(range(1)):
        img = cv2.imread(img_path)

        # 读取bounding box
        boxes = get_boxes(label_path)
        draw(boxes, img)

        if not save_dir.exists():
            save_dir.mkdir(parents=True)
        cv2.imwrite(Path(save_dir, img_path.name), img)